# remove_vectors.py
# 功能：清空 Chroma 数据库中指定集合的所有记录，或删除整个集合
# usage: python ./src/vector/remove_vectors.py

import os,sys
project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
sys.path.insert(0, project_root)
from config import project_config
import chromadb

CHROMA_DB_PATH = project_config.chroma_location  
COLLECTION_NAME = project_config.collection_name  

def delete_all_from_collection(db_path, collection_name):
    """连接 Chroma DB，清空指定集合中的所有记录"""
    try:
        client = chromadb.PersistentClient(path=db_path)
    except Exception as e:
        print(f"❌ 无法连接到 Chroma 数据库，请检查路径是否正确：{db_path}")
        print(f"错误信息：{e}")
        return

    try:
        collection = client.get_collection(name=collection_name)
    except ValueError:
        print(f"❌ 未找到集合 '{collection_name}'，请确认集合名称是否正确。")
        return
    except Exception as e:
        print(f"❌ 访问集合时发生错误：{e}")
        return

    # 获取当前记录数量
    count_before = collection.count()
    if count_before == 0:
        print(f"✅ 集合 '{collection_name}' 已经是空的，无需删除。")
        return

    print(f"📊 集合 '{collection_name}' 当前共有 {count_before} 条记录。")

    # 获取所有 ID（Chroma 删除必须提供 ID 列表）
    results = collection.get(include=[], limit=None)
    ids_to_delete = results['ids']

    try:
        collection.delete(ids=ids_to_delete)
        print(f"✅ 成功删除 {len(ids_to_delete)} 条记录。")
    except Exception as e:
        print(f"❌ 删除记录时出错：{e}")

    # 验证是否已清空
    count_after = collection.count()
    print(f"📊 删除后，集合中剩余记录数：{count_after}")

def delete_collection(db_path, collection_name):
    """（可选）直接删除整个集合"""
    try:
        client = chromadb.PersistentClient(path=db_path)
    except Exception as e:
        print(f"❌ 无法连接到 Chroma 数据库：{e}")
        return

    try:
        client.delete_collection(name=collection_name)
        print(f"🔥 集合 '{collection_name}' 已被彻底删除。")
    except ValueError:
        print(f"❌ 未找到集合 '{collection_name}'，无法删除。")
    except Exception as e:
        print(f"❌ 删除集合时出错：{e}")

if __name__ == '__main__':
    print(f"⚠️  正在准备清空集合：'{COLLECTION_NAME}'")

    # 方式一：清空集合内所有记录（保留集合结构）
    delete_all_from_collection(CHROMA_DB_PATH, COLLECTION_NAME)

    # 方式二：如果你希望直接删除整个集合（取消下面的注释即可）
    # print("\n" + "="*50)
    # print("⚠️  正在准备删除整个集合...")
    # delete_collection(CHROMA_DB_PATH, COLLECTION_NAME)